package advice.org.springframework.asm;

import advice.org.springframework.util.ClassUtilsAdvice;
import cn.hutool.core.collection.ConcurrentHashSet;
import com.dcits.agent.IAdvice;
import com.dcits.agent.utils.RedefineClassUtils;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.method.MethodDescription;
import net.bytebuddy.matcher.ElementMatcher;
import net.bytebuddy.matcher.ElementMatchers;
import org.apache.commons.lang3.StringUtils;
import org.kohsuke.MetaInfServices;
import org.springframework.asm.*;

import java.io.*;
import java.util.HashSet;
import java.util.Set;
import java.util.stream.Stream;

/**
 * @className: ClassReaderAdvice
 * @description:
 * @author: caoyangjie
 * @date: 2025/6/21
 **/
@MetaInfServices
public class ClassReaderAdvice implements IAdvice {

    public static Set<String> existClasses = new ConcurrentHashSet<>();
    public static Set<String> includeOtherClasses = new ConcurrentHashSet<>();
    @Advice.OnMethodEnter
    static void onMethodEnter(@Advice.Argument(value = 0,readOnly = false) byte[] classFileByte
    ) {
            classFileByte = modifyClassFile(classFileByte);
    }

    public static byte[] modifyClassFile(byte[] classBytes) {
        ClassWriter classWriter = new ClassWriter(0);
        SmartClassReader classReader = new SmartClassReader(classBytes, 0, classBytes.length);
        String className = classReader.getClassName();
        ClassVisitor classVisitor = new ClassReaderAdvice.ReplaceVisitor(classWriter, className);
        classReader.accept(classVisitor, 0);
        return classWriter.toByteArray();
    }

    /**
     * 加载增强类时,需要附带将增强类中创建的类也加载到类加载器
     */
    public static void loaderIncludeClasses() {
        includeOtherClasses.stream().forEach(className -> {
            if( !existClasses.contains(className) ) {
                existClasses.add(className);
                try{
                    ClassUtilsAdvice.replaceClass(className.replaceAll("\\/","\\."));
                } catch(Throwable e){
                }
            }
        });
        includeOtherClasses.clear();
    }

    @Override
    public void advice() {
        // 这里需要根据实际情况填写 redefine 方法
        ElementMatcher<? super MethodDescription> matcher = ElementMatchers
                .isConstructor()
                .and(ElementMatchers.takesArguments(byte[].class));
        // 增强当前类中所有的非构造函数的 public 方法
        RedefineClassUtils.redefineMethodWithAdvice(org.springframework.asm.ClassReader.class, matcher);
    }

    public static class SmartClassReader extends ClassReader {

        static String replaceSerlvetString(String value) {
            if( StringUtils.isEmpty(value)  ) {
                return value;
            }
            value = value.replaceAll("javax/servlet", "jakarta/servlet")
                            .replaceAll("javax\\.servlet", "jakarta\\.servlet");
            return value;
        }

        public SmartClassReader(byte[] classFileBuffer, int classFileOffset, int classFileLength) {
            super(classFileBuffer, classFileOffset, classFileLength);
        }

        @Override
        public String readUTF8(int offset, char[] charBuffer) {
            return replaceSerlvetString(super.readUTF8(offset, charBuffer));
        }

        @Override
        public String readClass(int offset, char[] charBuffer) {
            String className = replaceSerlvetString(super.readClass(offset, charBuffer));
            // 动态增强的类中引入的新类也需要使用 SmartClassLoader 来进行加载 否则会出现 同一个类在jvm中多分，会出现类型转换异常
//            if( StringUtils.isNotBlank(className) && className.startsWith("cn/com/zybank") ) {
//                includeOtherClasses.add(className);
//            }
            return className;
        }

        @Override
        public String readPackage(int offset, char[] charBuffer) {
            return replaceSerlvetString(super.readPackage(offset, charBuffer));
        }
    }

    public static class ReplaceVisitor extends ClassVisitor {
        private String className;
        public ReplaceVisitor(ClassVisitor classVisitor, String className) {
            super(Opcodes.ASM9, classVisitor);
            this.className = className;
        }

        static String[] replaceServletPackage(String[] classStrings) {
            if( classStrings==null || classStrings.length==0 ) {
                return classStrings;
            }
            String[] replaces = new String[classStrings.length];
            for( int i = 0; i < classStrings.length; i++ ) {
                String s = classStrings[i];
                if(StringUtils.isNotBlank(s)) {
                    replaces[i] = s.replaceAll("javax/servlet", "jakarta/servlet")
                            .replaceAll("javax\\.servlet", "jakarta\\.servlet");
                }
            }
            return replaces;
        }

        @Override
        public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
            // 替换类名
            superName = replaceServletPackage(new String[]{superName})[0];
            interfaces = replaceServletPackage(interfaces);
            super.visit(version, access, name.replace("javax/servlet", "jakarta/servlet"), signature, superName, interfaces);
        }

        @Override
        public AnnotationVisitor visitAnnotation(String descriptor, boolean visible) {
            descriptor = replaceServletPackage(new String[]{descriptor})[0];
            return super.visitAnnotation(descriptor, visible);
        }

        @Override
        public AnnotationVisitor visitTypeAnnotation(int typeRef, TypePath typePath, String descriptor, boolean visible) {
            descriptor = replaceServletPackage(new String[]{descriptor})[0];
            return super.visitTypeAnnotation(typeRef, typePath, descriptor, visible);
        }

        @Override
        public FieldVisitor visitField(int access, String name, String descriptor, String signature, Object value) {
            descriptor = replaceServletPackage(new String[]{descriptor})[0];
            return super.visitField(access, name, descriptor, signature, value);
        }

        @Override
        public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
            descriptor = replaceServletPackage(new String[]{descriptor})[0];
            exceptions= replaceServletPackage(exceptions);
            MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions);
            // 这里可以替换方法参数或返回值中的 javax.servlet 类型
            return new ClassReaderAdvice.ReplaceMethodVisitor(mv, this.className);
        }
        // 还可以在这里覆盖其他 visit 方法来替换更多类型引用


        @Override
        public void visitInnerClass(String name, String outerName, String innerName, int access) {
            super.visitInnerClass(name, outerName, innerName, access);
            if( StringUtils.isNotBlank(name) && name.startsWith("cn/com/zybank") ) {
                includeOtherClasses.add(name);
            }
        }
    }

    public static class ReplaceMethodVisitor extends MethodVisitor {
        private String className;
        public ReplaceMethodVisitor(MethodVisitor methodVisitor, String className) {
            super(Opcodes.ASM9, methodVisitor);
            this.className = className;
        }

        @Override
        public void visitTypeInsn(int opcode, String type) {
            // 替换类型
            if (type.startsWith("javax/servlet")) {
                type = type.replace("javax/servlet", "jakarta/servlet");
            }
            super.visitTypeInsn(opcode, type);
        }

        @Override
        public void visitFieldInsn(int opcode, String owner, String name, String descriptor) {
            // 替换字段的 owner 类型
            if (owner.startsWith("javax/servlet")) {
                owner = owner.replace("javax/servlet", "jakarta/servlet");
            }
            if( descriptor.contains("javax/servlet") ) {
                descriptor = descriptor.replaceAll("javax/servlet", "jakarta/servlet");
            }
            super.visitFieldInsn(opcode, owner, name, descriptor);
        }

        @Override
        public void visitMethodInsn(int opcode, String owner, String name, String descriptor, boolean isInterface) {
            // 替换方法的 owner 类型
            if (owner.startsWith("javax/servlet")) {
                owner = owner.replace("javax/servlet", "jakarta/servlet");
            }
            if( descriptor.contains("javax/servlet") ) {
                descriptor = descriptor.replaceAll("javax/servlet", "jakarta/servlet");
            }
            super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
        }

        @Override
        public void visitTryCatchBlock(Label start, Label end, Label handler, String type) {
            super.visitTryCatchBlock(start, end, handler, type);
        }

        @Override
        public void visitParameter(String name, int access) {
            super.visitParameter(name, access);
        }

        @Override
        public void visitLineNumber(int line, Label start) {
            super.visitLineNumber(line, start);
        }

        @Override
        public void visitCode() {
            super.visitCode();
        }
    }
}
